#include "icp.h"
#include "g2o_icp_edge.h"
#include "g2o_icp_vertex.h"

#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/solvers/dense/linear_solver_dense.h>

typedef g2o::BlockSolverX BlockSolverType;
typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;

void bundle_adjustment_g2o(const vector<Point3f> &pts1,
                           const vector<Point3f> &pts2,
                           Mat &R, Mat &t)
{
     // optimizer
     g2o::SparseOptimizer optimizer;
     optimizer.setVerbose(true);

     // solver
     auto solver =
          new g2o::OptimizationAlgorithmLevenberg(
               g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));
     optimizer.setAlgorithm(solver);

     // camera pose as the vertex
     VertexPose *pose = new VertexPose();
     pose->setId(0);
     pose->setEstimate(Sophus::SE3d());
     optimizer.addVertex(pose);

     // edges
     for (size_t i = 0; i < pts1.size(); i++)
     {
          EdgeProjectXYZRGBDPoseOnly *edge =
               new EdgeProjectXYZRGBDPoseOnly(Eigen::Vector3d(pts2[i].x, pts2[i].y, pts2[i].z));
          edge->setVertex(0, pose);
          edge->setMeasurement(Eigen::Vector3d(pts1[i].x, pts1[i].y, pts1[i].z));
          edge->setInformation(Eigen::Matrix3d::Identity());

          optimizer.addEdge(edge);
     }

     // do optimization
     chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
     optimizer.initializeOptimization();
     optimizer.optimize(10);
     chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
     chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
     cout << "optimization costs time: " << time_used.count() << " seconds." << endl;

     cout << endl
          << "after optimization:" << endl;
     cout << "T=\n"
          << pose->estimate().matrix() << endl;

     // convert to cv::Mat
     Eigen::Matrix3d R_ = pose->estimate().rotationMatrix();
     Eigen::Vector3d t_ = pose->estimate().translation();
     R = (Mat_<double>(3, 3) << R_(0, 0), R_(0, 1), R_(0, 2),
          R_(1, 0), R_(1, 1), R_(1, 2),
          R_(2, 0), R_(2, 1), R_(2, 2));
     t = (Mat_<double>(3, 1) << t_(0, 0), t_(1, 0), t_(2, 0));
}
